package com.example.springbootdemo.interceptor;

import com.google.common.util.concurrent.RateLimiter;
import org.springframework.http.HttpStatus;
import org.springframework.web.servlet.HandlerInterceptor;

import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * 限流拦截器
 */
public class RateLimitInterceptor implements HandlerInterceptor {

    private Map<String, RateLimiter> ipLimiter = new ConcurrentHashMap<>();

    @Override
    public boolean preHandle(HttpServletRequest request, HttpServletResponse response, Object handler) throws Exception {
        // java还提供获取代理、vpn....
        String ip = request.getRemoteHost();  // 最简单的获取IP的方式

        System.out.println("访问用户的IP: " + ip);
        RateLimiter rateLimiter = ipLimiter.get(ip); // 获取对应的ip的限流器

        if(null == rateLimiter) {
            rateLimiter = RateLimiter.create(2);
            ipLimiter.put(ip, rateLimiter);
        }

        boolean isAcquire = rateLimiter.tryAcquire(); //尝试获取令牌

        if(isAcquire) {
            return true;
        }else {
            response.setStatus(HttpStatus.TOO_MANY_REQUESTS.value());
            return false;
        }
    }


    /**
    // RateLimiter.create() 每秒钟产生1个令牌
    private RateLimiter rateLimiter = RateLimiter.create(1);

    @Override
    public boolean preHandle(HttpServletRequest request, HttpServletResponse response, Object handler) throws Exception {
        boolean isAcquire = rateLimiter.tryAcquire();  //尝试获取令牌
        if(isAcquire) {
            return true;
        }else {
            response.setStatus(HttpStatus.TOO_MANY_REQUESTS.value());
            return false;
        }
    }
    */
}
